import sys
sys.path.append("..")
sys.path.append(".")
sys.path.append('..')

import os
import torch
import torch.nn as nn
import torchvision
import torch.utils.data as Data
from torch.autograd import Variable
from torch.nn import functional as F
import utils.transforms as trans
import utils.utils as util
import layer.loss as ls
import utils.metric as mc
import shutil
from net.segnet import SiameseNet, FeatureResNet
from net.fusenet import fusenet
import cv2
import pickle
from config import cfg as cfg_deeplabv3_plus
from net.generateNet import generate_net
from net.backbone import build_backbone
import layer.loss as ls
import layer.function as fun
import numpy as np
import utils.utils as util
from PIL import Image, ImageDraw


import time
import os
import torchvision.transforms as transforms
import json

import train_fusenet

# 导入flask库的Flask类和request对象
from flask import Flask, render_template, request, jsonify, Response
app = Flask(__name__)
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

datasets = 'CD2014'
import cfgs.CD2014config as cfg
import dataset.CD2014 as dates


def check_dir(dir):
    if not os.path.exists(dir):
        #os.mkdir(dir)
        os.makedirs(dir)

def set_base_learning_rate_for_multi_layer(model):

    params_dict = dict(model.named_parameters())
    for layer_name, layer_param in params_dict.items():
        if 'embedding_layer' not in layer_name:
            if 'weight' in layer_name:
                yield layer_param

def set_2x_learning_rate_for_multi_layer(model):

    params_dict = dict(model.named_parameters())
    for layer_name, layer_param in params_dict.items():
        if 'embedding_layer' not in layer_name:
            if 'bias' in layer_name:
                yield layer_param

def set_10x_learning_rate_for_multi_layer(model):

    params_dict = dict(model.named_parameters())
    for layer_name, layer_param in params_dict.items():
        if 'embedding_layer' in layer_name:
            if 'weight' in layer_name:
                yield layer_param

def set_20x_learning_rate_for_multi_layer(model):

    params_dict = dict(model.named_parameters())
    for layer_name, layer_param in params_dict.items():
        if 'embedding_layer' in layer_name:
            if 'bias' in layer_name:
                yield layer_param

def untransform(transform_img,mean_vector):

    transform_img = transform_img.transpose(1,2,0)
    transform_img += mean_vector
    transform_img = transform_img.astype(np.uint8)
    transform_img = transform_img[:,:,::-1]
    return transform_img

def various_distance(out_vec_t0, out_vec_t1,dist_flag):
    if dist_flag == 'l2':
        distance = F.pairwise_distance(out_vec_t0, out_vec_t1, p=2)
    if dist_flag == 'l1':
        distance = F.pairwise_distance(out_vec_t0, out_vec_t1, p=1)
    if dist_flag == 'cos':
        distance = 1 - F.cosine_similarity(out_vec_t0, out_vec_t1)
    return distance

def normFUSE(d):
    ma = torch.max(d)
    mi = torch.min(d)
    dn = (d-mi)/(ma-mi)
    return dn

def process_points_to_roi(im_path , points_list , output_path) :
    # 读取图片名
    img = cv2.imread(im_path)

    # # 绘制mask
    zeros = np.zeros((img.shape) , dtype=np.uint8)

    for points in points_list :
        # 获取坐标信息,
        points = np.array([points] , dtype=np.int32)
        # 原本thickness = -1表示内部填充,这里不知道为什么会报错,只好不填充了
        cv2.polylines(zeros , points , isClosed=True , thickness=45 , color=(255 , 255 , 255))

    val_transform_det = trans.Compose([
        trans.Scale(cfg.TRANSFROM_SCALES) ,
    ])
    im = Image.fromarray(zeros)
    im = val_transform_det(im)
    format = 'png'
    save_name = im_path.split("/")[-1].replace("jpg" , "roi")
    print("im.size:", im.size)


    im.save(os.path.join(output_path , save_name) , format)
    # return os.path.join(output_path , 'ROI.bmp')
    return os.path.join(output_path , save_name)

def process_points_to_binary(im_path , points_list , output_path) :
    # 读取图片名
    img = cv2.imread(im_path)

    # 绘制mask
    zeros = np.zeros((img.shape) , dtype=np.uint8)
    im_in = Image.open(im_path)
    im = Image.fromarray(zeros)
    draw = ImageDraw.Draw(im)

    for points in points_list :
        print(points)
        point_new = list()
        for p in points :
            point_new.append((p[0] , p[1]))
        print(point_new)
        draw.polygon(point_new , fill=(1 , 1 , 1))

    del draw
    format = 'png'
    save_name = im_path.split("/")[-1].replace("jpg", "binary")

    im.save(os.path.join(output_path , save_name) , format)

    return os.path.join(output_path , save_name)

def flask_train(img1_path, img2_path, label_path):
    with open(label_path , "r") as f :
        data = json.load(f)

    points_list = []
    for ps in data["shapes"] :
        points_list.append(ps["points"])

    lable_dir_name = os.path.dirname(label_path)
    binary_file = process_points_to_binary(img2_path, points_list , lable_dir_name)
    roi_file = process_points_to_roi(img2_path, points_list , lable_dir_name)

    with open(train_fusenet.cfg.TRAIN_TXT_PATH , 'w') as f :
        f.write(img1_path + ' ')
        f.write(img2_path + ' ')
        f.write(binary_file + ' ')
        f.write(roi_file)

    train_fusenet.train();

def flask_validate(net, img1_path,img2_path,save_change_map_dir):
    net.eval()
    inputs1, input2, height,width = load_paired_images(img1_path,img2_path)

    inputs1, input2 = inputs1.cuda(), input2.cuda()
    inputs1, inputs2 = Variable(inputs1, volatile=True), Variable(input2, volatile=True)
    fuse, out_middle, out_bottom, out_final = net(inputs1, inputs2)

    fuse = normFUSE(fuse)
    fuse_image = np.uint8(255 * fuse.data.cpu().numpy()[0][0])
    print("fuse_image.shape:", fuse_image.shape)
    roi_dir_name = os.path.dirname(img2_path)
    print(os.path.join(roi_dir_name,img2_path.split("/")[-1].replace("jpg" , "roi")))
    if os.path.exists(os.path.join(roi_dir_name,img2_path.split("/")[-1].replace("jpg" , "roi"))):
        print("roi file exists!")
        roi = Image.open(os.path.join(roi_dir_name,img2_path.split("/")[-1].replace("jpg" , "roi")))
        roi = roi.convert("L")

        val_transform_det = trans.Compose([
            trans.Scale(cfg.TRANSFROM_SCALES) ,
        ])

        roi = val_transform_det(roi)

        roi = np.array(roi, dtype=np.uint8)
        print(np.unique(roi))
        print("roi.shape:", roi.shape)
        # fuse_image = roi * 255
        fuse_image *= (roi * 255)


    check_dir(save_change_map_dir)

    save_weight_fig_path = os.path.join(save_change_map_dir, 'output.jpg')

    cv2.imwrite(save_weight_fig_path, fuse_image)
    print(save_weight_fig_path)
    put_seg_mask(binary_img_path=save_weight_fig_path , img_path=img2_path, output_fold=save_change_map_dir)


def put_seg_mask(binary_img_path , img_path , output_fold) :
    # 读取图片
    binary_image = cv2.imread(binary_img_path)
    binary_image = np.where(binary_image > (binary_image.max() + binary_image.min()) / 2 , 255 , 0);
    cv2.imwrite(os.path.join(output_fold , 'binary_image.jpg') , binary_image)
    binary_image[: , : , 1] = 0
    binary_image[: , : , 2] = 0

    image = cv2.imread(img_path)
    image = cv2.resize(image , (binary_image.shape[0] , binary_image.shape[1]))


    # 画出mask
    zeros_mask = np.zeros((image.shape) , dtype=np.uint8)
    zeros_mask[: :1] = binary_image[: :1]
    print(zeros_mask.shape)

    try :
        # alpha 为第一张图片的透明度
        alpha = 1
        # beta 为第二张图片的透明度
        beta = 0.5
        gamma = 0
        # cv2.addWeighted 将原始图片与 mask 融合
        mask_img = cv2.addWeighted(image , alpha , zeros_mask , beta , gamma)
        cv2.imwrite(os.path.join(output_fold , 'mask_img.jpg') , mask_img)
    except :
        print('异常')

def data_transform(img1,img2):
    img1 = img1[:, :, ::-1]  # RGB -> BGR
    img1 = img1.astype(np.float64)
    img1 -= cfg.T0_MEAN_VALUE
    img1 = img1.transpose(2, 0, 1) # H×W×C -> C×H×W
    img1 = torch.from_numpy(img1).float()
    img2 = img2[:, :, ::-1]  # RGB -> BGR
    img2 = img2.astype(np.float64)
    img2 -= cfg.T1_MEAN_VALUE
    img2 = img2.transpose(2, 0, 1)
    img1 = img1[np.newaxis, :]
    img2 = img2[np.newaxis, :]
    img2 = torch.from_numpy(img2).float()

    return img1,img2

def load_paired_images(img1_path,img2_path):
    ####### load images #############
    img1 = Image.open(img1_path).convert('RGB')
    img2 = Image.open(img2_path).convert('RGB')

    #print filename
    height,width,_ = np.array(img1,dtype= np.uint8).shape

    val_transform_det = trans.Compose([
        trans.Scale(cfg.TRANSFROM_SCALES) ,
    ])
    img1 = val_transform_det(img1)
    img2 = val_transform_det(img2)

    img1 = np.array(img1 , dtype=np.uint8)
    img2 = np.array(img2 , dtype=np.uint8)

    img1,img2 = data_transform(img1, img2)
    # img1 = np.array(img1,dtype= np.uint8)
    # img2 = np.array(img2,dtype= np.uint8)

    return img1, img2, int(height),int(width)


def validate_by_images(img1_path , img2_path) :
    torch.backends.cudnn.benchmark = True

    torch.cuda.empty_cache()

    val_transform_det = trans.Compose([
        trans.Scale(cfg.TRANSFROM_SCALES) ,
    ])
    val_data = dates.Dataset(cfg.VAL_DATA_PATH , cfg.VAL_LABEL_PATH ,
                             cfg.VAL_TXT_PATH , 'val' , transform=True ,
                             transform_med=val_transform_det)
    val_loader = Data.DataLoader(val_data , batch_size=1 ,
                                 shuffle=False , num_workers=0 , pin_memory=True , drop_last=True)
    ######  build  models ########
    pretrained_net = FeatureResNet()
    Siamese = SiameseNet(pretrained_net)
    model = fusenet(Siamese)
    checkpoint_path = cfg.TRAINED_BEST_PERFORMANCE_CKPT

    model.load_state_dict(
        {k.replace('module.' , '') : v for k , v in torch.load(checkpoint_path)['state_dict'].items()})
    print('validate begin.')

    model = model.cuda()
    SAVE_PATH = './data'
    output_dir = os.path.join(SAVE_PATH , 'output')
    check_dir(SAVE_PATH)
    check_dir(output_dir)

    flask_validate(model , img1_path , img2_path , output_dir)
    print('Processing completed!')

#------------------------------------------------------1.加载模型--------------------------------------------------------------
torch.backends.cudnn.benchmark = True
torch.cuda.empty_cache()
pretrained_net = FeatureResNet()
Siamese = SiameseNet(pretrained_net)
model = fusenet(Siamese)
checkpoint_path = cfg.TRAINED_BEST_PERFORMANCE_CKPT
model.load_state_dict({k.replace('module.', ''): v for k, v in torch.load(checkpoint_path)['state_dict'].items()})
# print('module load begin.')

model = model.cuda()
SAVE_PATH = './static'
output_dir = os.path.join(SAVE_PATH,'output')
check_dir(SAVE_PATH)
check_dir(output_dir)

img1_path , img2_path = './data/test/in000474.jpg', './data/test/in001700.jpg'
# flask_validate(model, img1_path, img2_path, output_dir)
print('module load completed!')

#------------------------------------------------------2.服务返回--------------------------------------------------------------
# 定义回调函数，接收来自/的post请求，并返回预测结果
@app.route("/send_image", methods=['POST'])
def return_result():
    received_file01 = request.files['img01']
    received_file02 = request.files['img02']
    imageFileName01 = received_file01.filename
    imageFileName02 = received_file02.filename
    if save_received_file(received_file01 , imageFileName01) and save_received_file(received_file02 , imageFileName02) :
        received_dirPath = './static/received_images'
        img1_path = os.path.join(received_dirPath , imageFileName01)
        img2_path = os.path.join(received_dirPath , imageFileName02)
        print(img1_path , img2_path)
        flask_validate(model , img1_path , img2_path , output_dir)

        # 创建response对象
        image = open("./static/output/mask_img.jpg", 'rb')
        resp = Response(image , mimetype="image/jpeg")
        return resp
    else :
        return 'failed'

# 访问首页时的调用函数
@app.route('/')
def index_page():
    return render_template('index.html')

def save_received_file(received_file, imageFileName):
    startTime = time.time()
    if received_file:
        received_dirPath =  './static/received_images'
        if not os.path.isdir(received_dirPath):
            os.makedirs(received_dirPath)
        imageFilePath = os.path.join(received_dirPath, imageFileName)
        received_file.save(imageFilePath)
        print('image file saved to %s' % imageFilePath)
        usedTime = time.time() - startTime
        print('接收图片并保存，总共耗时%.2f秒' % usedTime)

        return True
    else:
        return False

# 使用predict_image这个API服务时的调用函数
@app.route("/upload_image", methods=['POST'])
def anyname_you_like():
    startTime = time.time()
    received_file01 = request.files['input_image01']
    received_file02 = request.files['input_image02']
    imageFileName01 = received_file01.filename
    imageFileName02 = received_file02.filename
    if save_received_file(received_file01, imageFileName01) and save_received_file(received_file02, imageFileName02):
        received_dirPath =  './static/received_images'
        img1_path = os.path.join(received_dirPath , imageFileName01)
        img2_path = os.path.join(received_dirPath , imageFileName02)
        print(img1_path, img2_path)
        flask_validate(model , img1_path , img2_path , output_dir)

        return render_template("result.html")
    else:
        return 'failed'


# 使用train_image这个API服务时,训练函数
@app.route("/train_image", methods=['POST'])
def train_images():
    received_file01 = request.files['img01']
    received_file02 = request.files['img02']
    received_file03 = request.files['label']
    imageFileName01 = received_file01.filename
    imageFileName02 = received_file02.filename
    labelFileName = received_file03.filename
    if save_received_file(received_file01 , imageFileName01) and save_received_file(received_file02 , imageFileName02)\
            and save_received_file(received_file03, labelFileName):
        current_dir = os.path.dirname(os.path.abspath(__file__))
        print(current_dir)
        received_dirPath = os.path.join(current_dir , 'static/received_images')
        current_dir = os.path.dirname(os.path.abspath(__file__))
        print(current_dir)
        img1_path = os.path.join(received_dirPath , imageFileName01)
        img2_path = os.path.join(received_dirPath , imageFileName02)
        label_path = os.path.join(received_dirPath , labelFileName)
        print(img1_path , img2_path)

        flask_train(img1_path, img2_path, label_path);

        checkpoint_path = cfg.TRAINED_BEST_PERFORMANCE_CKPT
        model.load_state_dict({k.replace('module.' , '') : v for k , v in torch.load(checkpoint_path)['state_dict'].items()})
        flask_validate(model , img1_path , img2_path , output_dir)

        # 创建response对象
        image = open("./static/output/mask_img.jpg" , 'rb')
        resp = Response(image , mimetype="image/jpeg")
        return resp
    else :
        return 'failed'


# 主函数
if __name__ == "__main__":
    app.run("127.0.0.1", port=5000)